import gradio as gr
import joblib
import numpy as np
from PIL import Image

# 加载模型
model = joblib.load('LGBMClassifier_best_model.pkl')

# 定义一个函数来调整图片尺寸并预测
def predict_image(image):
    # 确保image是PIL Image对象
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    # 调整图片尺寸以匹配模型的输入形状，这里假设模型期望的输入形状是(32, 32)，你需要根据你的模型调整这个尺寸
    img = image.resize((32, 32))
    # 将图片转换为数组
    image_array = np.array(img)
    # 展平图片数组
    image_array = image_array.flatten()
    # 如果展平后的数组长度不等于模型期望的特征数量，你需要进行进一步的特征提取或降维
    # 这里假设模型期望1024个特征，所以我们需要将展平后的数组裁剪或填充到1024个特征
    if len(image_array) > 1024:
        image_array = image_array[:1024]  # 裁剪到1024个特征
    elif len(image_array) < 1024:
        # 如果特征数量少于1024，你需要填充数组或使用其他方法来增加特征数量
        # 这里只是一个示例，实际应用中可能需要更复杂的处理
        image_array = np.pad(image_array, (0, 1024 - len(image_array)), 'constant')
    # 将图片数组重塑为模型期望的形状
    image_array = image_array.reshape(1, -1)
    # 预测
    prediction = model.predict(image_array)
    # 返回预测结果
    return "猫" if prediction[0] == 0 else "狗"

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),  # 确保输入类型为PIL Image
    outputs="text"
)

# 启动应用，并设置share=True以创建公共链接
iface.launch(share=True)